package org.multi.core;

import java.util.Iterator;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.BiPredicate;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;
import org.multi.core.function.TripleConsumer;
import org.multi.core.function.TripleIntConsumer;
import org.multi.core.function.TripleIntFunction;
import org.multi.core.state.IntState;

public interface BiMulti<T, U> {

    static void stop() {
        StopException.stop();
    }

    static <T, U> BiMulti<T, U> empty() {
        return consumer -> {
        };
    }

    static <T, U> BiMulti<T, U> create(Consumer<BiConsumer<T, U>> action) {
        return action::accept;
    }

    static <T, U> BiMulti<T, U> zip(Multi<T> multi, Iterable<U> iterable) {
        return consumer -> {
            Iterator<U> iterator = iterable.iterator();
            multi.consumeOn(item -> {
                if (iterator.hasNext()) {
                    consumer.accept(item, iterator.next());
                }
                stop();
            });
        };
    }

    default <E> BiMulti<E, U> mapKey(Function<T, E> mapper) {
        return consumer -> consume((key, value) -> {
            consumer.accept(mapper.apply(key), value);
        });
    }

    default <E> BiMulti<T, E> mapValue(Function<U, E> mapper) {
        return consumer -> consume((key, value) -> {
            consumer.accept(key, mapper.apply(value));
        });
    }

    default <E, R> BiMulti<E, R> map(Function<T, E> mapperKey, Function<U, R> mapperValue) {
        return consumer -> consume((key, value) -> {
            consumer.accept(mapperKey.apply(key), mapperValue.apply(value));
        });
    }

    default <R> Multi<R> map(BiFunction<T, U, R> mapper) {
        return consumer -> consume((key, value) -> {
            consumer.accept(mapper.apply(key, value));
        });
    }

    default BiMulti<T, U> filterKey(Predicate<T> test) {
        return consumer -> consume((key, value) -> {
            if (test.test(key)) {
                consumer.accept(key, value);
            }
        });
    }

    default BiMulti<T, U> filterValue(Predicate<U> test) {
        return consumer -> consume((key, value) -> {
            if (test.test(value)) {
                consumer.accept(key, value);
            }
        });
    }

    default BiMulti<T, U> filter(BiPredicate<T, U> test) {
        return consumer -> consume((key, value) -> {
            if (test.test(key, value)) {
                consumer.accept(key, value);
            }
        });
    }

    void consume(BiConsumer<T, U> consumer);

    default BiMulti<T, U> onFailure(Consumer<Throwable> onFailure) {
        return consumer -> consume(consumer, onFailure);
    }

    default void consume(BiConsumer<T, U> consumer, Consumer<Throwable> onFailure) {
        try {
            consume(consumer);
        } catch (Throwable error) {
            onFailure.accept(error);
        }
    }

    default void consume(int from, BiConsumer<T, U> headConsumer, int to, BiConsumer<T, U> otherConsumer) {
        if (from < 0 || to < 0 || from > to) {
            consume(headConsumer);
        }
        IntState state = new IntState(from);
        consume((key, value) -> {
            if (state.value < to) {
                state.value++;
                headConsumer.accept(key, value);
            } else {
                otherConsumer.accept(key, value);
            }
        });
    }

    default void consume(BiConsumer<T, U> headConsumer, int to, BiConsumer<T, U> otherConsumer) {
        consume(0, headConsumer, to, otherConsumer);
    }

    default void consumeOn(BiConsumer<T, U> consumer) {
        try {
            consume(consumer);
        } catch (StopException ignored) {

        }
    }

    default void consumeWithIndex(TripleIntConsumer<T, U> consumer) {
        IntState state = new IntState(0);
        consume((key, value) -> consumer.accept(state.value++, key, value));
    }

    default void forEach(BiConsumer<T, U> consumer) {
        consume(consumer);
    }

    default <E, R> BiMulti<E, R> mapMulti(TripleConsumer<T, U, BiConsumer<E, R>> function) {
        return consumer -> consume((key, value) -> function.accept(key, value, consumer));
    }

    default <R> Multi<R> mapIndexed(TripleIntFunction<T, U, R> function) {
        return consumer -> consumeWithIndex((index, key, value) -> consumer.accept(function.apply(index, key, value)));
    }

    default <E, R> BiMulti<E, R> mapIndexed(TripleIntFunction<T, U, E> keyFunction, TripleIntFunction<T, U, R> valueFunction) {
        return consumer -> consumeWithIndex((index, key, value) -> {
            consumer.accept(keyFunction.apply(index, key, value), valueFunction.apply(index, key, value));
        });
    }

    default boolean isEmpty() {
        boolean[] state = {true};
        consumeOn((key, value) -> {
            state[0] = false;
            stop();
        });
        return state[0];
    }

    default BiMulti<T, U> flatmap(BiFunction<T, U, BiMulti<T, U>> mapper) {
        return consumer -> consume((key, value) -> {
            mapper.apply(key, value).consume(consumer);
        });
    }

    default BiMulti<T, U> take(long count) {
        return consumer -> {
            long[] state = {count};
            consumeOn((key, value) -> {
                consumer.accept(key, value);
                if (--state[0] == 0) {
                    stop();
                }
            });
        };
    }

    default BiMulti<T, U> takeWhile(BiPredicate<T, U> test) {
        return consumer -> {
            consumeOn((key, value) -> {
                if (test.test(key, value)) {
                    consumer.accept(key, value);
                } else {
                    stop();
                }
            });
        };
    }

    default BiMulti<T, U> drop(long count) {
        return consumer -> {
            long[] state = {count};
            consume((key, value) -> {
                if (state[0] > 0) {
                    state[0]--;
                } else {
                    consumer.accept(key, value);
                }
            });
        };
    }

    default BiMulti<T, U> dropWhile(BiPredicate<T, U> test) {
        return consumer -> {
            boolean[] state = {true};
            consume((key, value) -> {
                if (state[0] && test.test(key, value)) {
                    return;
                }
                state[0] = false;
                consumer.accept(key, value);
            });
        };
    }

    default BiMulti<T, U> peek(BiConsumer<T, U> action) {
        return consumer -> consume(consumer.andThen(action));
    }
}
